import pandas as pd
import hashlib
import pefile
import os
import shutil

class DataCleaner:
    def __init__(self):
        self.seen_hashes = set()
    
    def _calculate_hash(self, file_path):
        """计算文件哈希值用于去重"""
        with open(file_path, 'rb') as f:
            return hashlib.sha256(f.read()).hexdigest()
    
    def is_valid_pe(self, file_path):
        """验证PE文件有效性"""
        try:
            pe = pefile.PE(file_path)
            return pe.is_valid()
        except:
            return False
    
    def clean_dataset(self, raw_dir, clean_dir):
        """执行清洗流程"""
        os.makedirs(clean_dir, exist_ok=True)
        
        for file_name in os.listdir(raw_dir):
            file_path = os.path.join(raw_dir, file_name)
            
            # 有效性检查
            if not self.is_valid_pe(file_path):
                continue
                
            # 去重检查
            file_hash = self._calculate_hash(file_path)
            if file_hash in self.seen_hashes:
                continue
            self.seen_hashes.add(file_hash)
            
            # 保存清理后的文件
            shutil.copy(file_path, os.path.join(clean_dir, file_name)) 